using System;
using System.Reflection;
using System.Collections;
using NBody;
using NBody.Analysis;
using OptionParser;
using DataFile;

[assembly: AssemblyTitle ("NBodyVelocity")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Compute various velocity quantities.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class VelocityOptions : InputOutputGetOptions
{
    [Option ("-b", "Number of bins")]
    public int NumBins;
    
    [Option ("-spacing", "Bin spacing")]
    public Bin.Spacing Spacing;
    
    [Option ("-r", "Output radial velocity versus radius (10,000 points only)")]
    public bool RadialVelocity;
    
    [Option ("-s", "Output log(r), v")]
    public bool SemiLogplot;
    
    [Option ("-l", "Output log(r), log(v)")]
    public bool Logplot;
    
    public VelocityOptions()
    {
        NumBins = 100;
        Spacing = Bin.Spacing.Logarithmic;
    }
}

public class NBodyVelocity
{
    public static void Main(string[] args)
    {
        VelocityOptions opts = new VelocityOptions();
        opts.ProcessArgs(args, ".vel");
                        
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        DoSphericalDispersions(s);
        
        Table data;
        if (opts.RadialVelocity)
        {
            DoRadialVelocity(s, 10000, out data);
            data.Print(opts.OutFile);
        }
        else
        {
            VelocityBin bin;
            data = new Table(opts.NumBins, 3);
            data.Clear();
            bin = new VelocityBin();
            data.Description = "Radius, Radial velocity, Tangential velocity";
            
            
            bin.NumBins = opts.NumBins;
            bin.BinSpacing = opts.Spacing;
            
            bin.VelocityDirection = VelocityBin.Direction.Radial;
            bin.Analyse(s);
            data.SetColumn(0, bin.Data.Column(0));
            data.SetColumn(1, bin.Data.Column(1));
            
            bin.VelocityDirection = VelocityBin.Direction.Tangential;
            bin.Analyse(s);
            data.SetColumn(2, bin.Data.Column(1));
            
            if (opts.SemiLogplot)
                data.Print(opts.OutFile, new DataFile.Transform[] { Math.Log10, DataFile.Table.IdentityTransform, DataFile.Table.IdentityTransform } );
            else if (opts.Logplot)
                data.Print(opts.OutFile, new DataFile.Transform(Math.Log10));
            else
                data.Print(opts.OutFile);
        }
    }
    
    public static void DoSphericalDispersions(NBodySystem s)
    {
        double vr_bar = 0, vt_bar = 0;
        foreach (Particle p in s)
        {
            double tmp = p.RadialVelocity;
            vr_bar += tmp;
            tmp = p.TangentialVelocity;
            vt_bar += tmp;
        }
        vr_bar /= (double)s.NumParts;
        vt_bar /= (double)s.NumParts;
        
        double vr_disp = 0, vt_disp = 0;
        foreach (Particle p in s)
        {
            double tmp = p.RadialVelocity;
            vr_disp += (tmp - vr_bar) * (tmp - vr_bar);
            tmp = p.TangentialVelocity;
            vt_disp += (tmp - vt_bar) * (tmp - vt_bar);
        }
        vr_disp /= (double)s.NumParts;
        vt_disp /= (double)s.NumParts;
        
        Console.WriteLine("vr2_disp = {0}\nvt2_disp = {1}", vr_disp, vt_disp);
    }
    
    public static void DoCylindricalDispersions(NBodySystem s)
    {
        double vr_bar = 0, vz_bar = 0;
        foreach (Particle p in s)
        {
            double tmp = p.CylindricalRadialVelocity;
            vr_bar += tmp;
            tmp = p.Vz;
            vz_bar += tmp;
        }
        vr_bar /= (double)s.NumParts;
        vz_bar /= (double)s.NumParts;
        
        double vr_disp = 0, vz_disp = 0;
        foreach (Particle p in s)
        {
            double tmp = p.CylindricalRadialVelocity;
            vr_disp += (tmp - vr_bar) * (tmp - vr_bar);
            tmp = p.Vz;
            vz_disp += (tmp - vz_bar) * (tmp - vz_bar);
        }
        vr_disp /= (double)s.NumParts;
        vz_disp /= (double)s.NumParts;
        
        Console.WriteLine("vr2_disp = {0}\nvz2_disp = {1}", vr_disp, vz_disp);
    }
    
    public static void DoRadialVelocity(NBodySystem s, int n, out Table data)
    {
        if (n == 0)
            n = s.NumParts;
    
        data = new Table(n, 2, "Log10(Radius), Velocity");
        int c = 0;
        int mod = (int)((double)s.NumParts / (double)n);
        for (int i = 0; i < s.NumParts; i++)
        {
            if (i % mod == 0)
            {
                data[c, 0] = Math.Log10(s[i].Radius);
                data[c, 1] = s[i].RadialVelocity;
                c++;
                if (c == n)
                    break;
            }
        }
    }
    
    public static void DoVelocitySlice(NBodySystem s, double radius, double width, int n, out Table data)
    {
        // collect all particles in this range
        ArrayList list = new ArrayList();
        double vmin = 0.0, vmax = 0.0;
        foreach (Particle p in s)
        {
            double r = p.Radius;
            if (r > radius - 0.5 * width && r < radius + 0.5 * width)
            {
                list.Add(p);
                double vr = p.RadialVelocity;
                if (vr > vmax)
                    vmax = vr;
                else if (vr < vmin)
                    vmin = vr;
            }
        }
        vmax += 1e-10;
        Console.Error.WriteLine("Vmax = {0}, Vmin = {1}", vmax, vmin);
        
        double deltav = (vmax - vmin) / (double)n;
        data = new Table(n, 2);
        foreach (Particle p in list)
        {
            int pos = (int)((p.RadialVelocity - vmin) / deltav);
            data[pos, 1] += 1.0;
        }
        for (int i = 0; i < n; i++)
            data[i, 0] = vmin + 0.5 * deltav + deltav * i;
    }
}
